# -*- coding: utf-8 -*-
# @Time    : 2022/7/24 19:54
# @Author  : Aweo0419


import torch
from model_select import select
import cv2
from torch.autograd import Variable
from torchvision import transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def img_read(img):
    # img = cv2.imread(img, cv2.IMREAD_COLOR)
    img = cv2.resize(img, (32, 32))
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    img = Variable(torch.unsqueeze(transform(img), dim=0).float(), requires_grad=False)
    return img

def model_converter():
    # model = torch.load('CNN5v2.pth').to(device)  # 这里保存的是完整模型
    # model.eval()
    model_path = 'CNN5v2.pth'
    model = select('CNN5')
    # model = SEnet_train.CNN(n_class=2)
    model.load_state_dict(torch.load(model_path))

    t_img = cv2.imread("img_data/nag_img/1.png")
    test_ar = img_read(t_img)
    input_names = ['input']
    output_names = ['output']
    torch.onnx.export(model, test_ar, 'test.onnx',
                      export_params=True,
                      verbose=True,
                      input_names=input_names,
                      output_names=output_names)



if __name__ == '__main__':
    model_converter()